/**
 * @Author yang
 * @Description mongodb 工具类， 包括增删改查方法
 * @Todo 1. 采用乐观锁 控制事务
 * @Date 2020/12/25
 * @Time 8:58
 **/
package mongo_go_driver_demo

import (
	"context"
	"errors"
	"fmt"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
	"log"
	"time"
)

const (
	RETRY_COUNT    = 1   // 客户端连接失败后的重试 次数
	OPTION_TIMEOUT = 100 // 增删改查在客户端的超时时间, 单位:秒
)

var (
	// 初始化配置信息
	hosts = "39.104.81.10:30000,39.104.50.203:30000,39.104.17.200:30000"

	// 客户端
	client *mongo.Client
)

// 初始化连接
func init() {
	if ping(RETRY_COUNT) {
		log.Printf("==> mongdb 连接成功!!")
	}
}

// 执行连接操作
func connect() {
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	// mongodb://[username:password@]host1[:port1][,...hostN[:portN]][/[defaultauthdb][?options]]
	clientCnn, err := mongo.Connect(ctx, options.Client().ApplyURI(fmt.Sprintf("mongodb://%v", hosts)))
	if err != nil {
		log.Println("=====> ", err)
	}
	client = clientCnn
}

// ping 测试连通性, retry 指定重试次数, 当前连接不可用时, 自动重连, 默认为 1 次
func ping(retry int) (ok bool) {
	if client == nil {
		connect()
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	err := client.Ping(ctx, readpref.Primary())
	if err != nil {
		log.Printf("===>> 服务器连接中断, 错误消息为: %v \n", err)
		if retry > 0 {
			connect()
			return ping(retry - 1)
		}
		return
	}
	ok = true
	return
}

// Insert 新增, 返回 主键 和 异常, https://docs.mongodb.com/manual/tutorial/insert-documents/
func Insert(database, collection string, document interface{}, opts ...*options.InsertOneOptions) (id string, err error) {
	if !ping(RETRY_COUNT) {
		return "", errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	res, err := coll.InsertOne(ctx, document, opts...)
	if err != nil {
		return
	}
	id = ObjectId2String1(res.InsertedID)
	// 添加创建时间 / 更新时间
	resCount, err := updateCUTime(database, collection, id)
	if resCount.ModifiedCount < 1 {
		err = errors.New(fmt.Sprintf("mongodb 插入数据添加 创建时间/更新时间失败, 数据库索引id为: %v, 数据信息为: %v", id, document))
	}
	return
}

// InsertWithIndex 带索引插入 todo
func InsertWithIndex(database, collection string, document interface{}, uniqueFields ...string) (id string, err error) {
	if !ping(RETRY_COUNT) {
		return "", errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	index, err := addUniqueIndex(&ctx, coll, uniqueFields...)
	if err != nil {
		return
	}
	if len(index) < 1 {
		err = errors.New(fmt.Sprintf("唯一索引创建失败, 要索引的信息为: %v-%v-%v", database, collection, uniqueFields))
		return
	}
	res, err := coll.InsertOne(ctx, document)
	if err != nil {
		return
	}
	id = ObjectId2String1(res.InsertedID)
	// 添加创建时间 / 更新时间
	resCount, err := updateCUTime(database, collection, id)
	if resCount.ModifiedCount < 1 {
		err = errors.New(fmt.Sprintf("mongodb 插入数据添加 创建时间/更新时间失败, 数据库索引id为: %v, 数据信息为: %v", id, document))
	}
	return
}

// 为字段添加唯一索引, 传入多个字段表示创建 联合索引, 默认所有索引为正序
func addUniqueIndex(ctx *context.Context, collection *mongo.Collection, fields ...string) ([]string, error) {
	models := []mongo.IndexModel{}
	uniqueFlag := true
	//var expiredSeconds int32 = 1000
	bsonDFields := bson.D{}
	for _, v := range fields {
		bsonDFields = append(bsonDFields, bson.E{v, 1})
	}
	model := mongo.IndexModel{
		Keys: bsonDFields,
		Options: &options.IndexOptions{
			Unique: &uniqueFlag,
		},
	}
	models = append(models, model)
	return addManyIndex(ctx, collection, &models)
}

// dropOneIndex 删除单个索引
func dropOneIndex(ctx *context.Context, collection *mongo.Collection, indexName string) error {
	info, err := collection.Indexes().DropOne(*ctx, indexName)
	log.Printf("===>> mongodb 删除索引[%v]--[%v]--[%v]的执行结果信息为: %v\n", collection.Database().Name(), collection.Name(), indexName, info)
	return err
}

// dropIndexes 删除多个索引
func dropAllIndex(ctx *context.Context, collection *mongo.Collection) error {
	info, err := collection.Indexes().DropAll(*ctx)
	log.Printf("===>> mongodb 删除索引[%v]--[%v]的执行结果信息为: %v\n", collection.Database().Name(), collection.Name(), info)
	return err
}

// 为字段添加多个索引
func addManyIndex(ctx *context.Context, collection *mongo.Collection, models *[]mongo.IndexModel, opts ...*options.CreateIndexesOptions) ([]string, error) {
	return collection.Indexes().CreateMany(*ctx, *models, opts...)
}

// indexExists 查看指定索引是否存在
func indexExists(ctx *context.Context, collection *mongo.Collection, indexName string) (err error) {
	list, err := collection.Indexes().List(*ctx)
	if err != nil {
		return
	}
	var indexes bson.A
	err = list.All(*ctx, &indexes)
	if err != nil {
		return
	}
	ok := false
	for _, v := range indexes {
		fmt.Println("=====> ", v)
		d := v.(bson.D)
		data := d[len(d)-1:][0]
		ok = data.Value == indexName
	}
	if !ok {
		err = errors.New("要查询的索引不存在")
	}
	return
}

// FindOne 获取符合给定条件的第一个结果, https://docs.mongodb.com/manual/tutorial/query-documents/
// res 为 *XXX 其中 XXX 为目标结构体
func FindOne(database, collection string, filter, result interface{}) (err error) {
	if !ping(RETRY_COUNT) {
		return errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	err = coll.FindOne(ctx, filter).Decode(result)
	return
}

// Find 条件查询, https://docs.mongodb.com/manual/tutorial/query-documents/
// res 为 *[]XXX 其中 XXX 为目标结构体
func Find(database, collection string, filter, result interface{}) (err error) {
	if !ping(RETRY_COUNT) {
		return errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	cursor, err := coll.Find(ctx, filter)
	if err != nil {
		return err
	}
	err = cursor.All(ctx, result)
	return
}

// Aggregate 聚合查询, https://docs.mongodb.com/manual/reference/sql-aggregation-comparison/
// res 为 *[]XXX 其中 XXX 为目标对象(当使用 groupby 时， 需要自定义传入数据结构）
func Aggregate(database, collection string, filter, result interface{}) (err error) {
	if !ping(RETRY_COUNT) {
		return errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	cursor, err := coll.Aggregate(ctx, filter)
	if err != nil {
		return err
	}
	err = cursor.All(ctx, result)
	return
}

// UpdateOne 条件更新, https://docs.mongodb.com/manual/tutorial/update-documents/
// 说明：分为 局部字段更新（ModifiedCount） / 整体文档覆盖（UpsertedCount） 两种类型, 生效条件详见链接文档
func UpdateOne(database, collection string, filter, updateExpression interface{}, opts ...*options.UpdateOptions) (resCount *mongo.UpdateResult, err error) {
	if !ping(RETRY_COUNT) {
		return nil, errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	resCount, err = coll.UpdateOne(ctx, filter, updateExpression, opts...)
	return
}

// UpdateOne 插入数据后, 添加插入/修改时间为当前时间
func updateCUTime(database, collection string, mongoId string, opts ...*options.UpdateOptions) (resCount *mongo.UpdateResult, err error) {
	if !ping(RETRY_COUNT) {
		return nil, errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	hex, err := primitive.ObjectIDFromHex(mongoId)
	if err != nil {
		err = errors.New(fmt.Sprintf("mongodb [%v]--[%v]--[%v] 索引异常!", database, collection, mongoId))
		return
	}
	filter := bson.D{
		{"_id", bson.D{
			{"$eq", hex},
		}},
	}
	updateExpression := bson.D{
		{"$currentDate", bson.D{
			{"c_time", true},
			{"u_time", true},
		}},
	}
	resCount, err = coll.UpdateOne(ctx, filter, updateExpression, opts...)
	return
}

// UpdateMany 条件更新, https://docs.mongodb.com/manual/tutorial/update-documents/
// 说明：分为 局部字段更新（ModifiedCount） / 整体文档覆盖（UpsertedCount） 两种类型, 生效条件详见链接文档
func UpdateMany(database, collection string, filter, updateExpression interface{}, opts ...*options.UpdateOptions) (resCount *mongo.UpdateResult, err error) {
	if !ping(RETRY_COUNT) {
		return nil, errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	resCount, err = coll.UpdateMany(ctx, filter, updateExpression, opts...)
	return
}

// DeleteOne 删除符合条件的第一个文本, 若没有异常且结果大于0, 才是删除成功
func DeleteOne(database, collection string, filter interface{}) (count int64, err error) {
	if !ping(RETRY_COUNT) {
		return 0, errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	countNum, err := coll.DeleteOne(ctx, filter)
	if countNum != nil {
		count = countNum.DeletedCount
	}

	fmt.Println("=====> count", count, err)
	return
}

// DeleteMany 删除符合条件的所有文档,
// !!! 谨慎使用, 若传参条件为空, 则可能导致删除整个数据库
func DeleteMany(database, collection string, filter interface{}) (count int64, err error) {
	if !ping(RETRY_COUNT) {
		return 0, errors.New("mongdb 连接失败!!")
	}
	ctx, cancel := context.WithTimeout(context.Background(), OPTION_TIMEOUT*time.Second)
	defer cancel()
	coll := client.Database(database).Collection(collection)
	countNum, err := coll.DeleteMany(ctx, filter)
	if countNum != nil {
		count = countNum.DeletedCount
	}
	fmt.Println("=====> DeleteMany", countNum, err)
	return
}

// ============================================其他工具类=========================================
// ObjectId2String primitive.ObjectID(mongodb 数据库主键类型) 转 string
func ObjectId2String(id interface{}) (idStr string, err error) {
	if id == nil {
		return "", errors.New("参数不能为空!")
	}
	if objectID, ok := id.(primitive.ObjectID); ok {
		return fmt.Sprintf("%v", objectID.Hex()), nil
	} else {
		return "", errors.New("参数类型不是 primitive.ObjectID, 请核对后重试!")
	}
}

func ObjectId2String1(id interface{}) string {
	id2String, _ := ObjectId2String(id)
	return id2String
}
